import os
import time
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import n4j.call_liantong as nclt
import math
import yinlichili.gravity as gra
import yinlichili.yuchuli as yuc
import yinlichili.loss_function as lossf


def main(length=1, gravity=1):
    # a = nclt.main(1)
    # # weiz = pd.read_excel('D:\BDZ.xlsx')
    # weiz = nclt.main(0)
    yuchuli_len = 30 - (length - 1) / 2
    if gravity == 1:
        Gravity = 200
    else:
        Gravity = 10
    aa = nclt.main(1)
    weizz = nclt.main(0)
    for first_len in range(length):
        Loss_out = {}
        for i_gravity in range(int((gravity - 1) / 10)):
            i_gravity = i_gravity * 10
            env_value_name = '初始长' + str(yuchuli_len + first_len) + '重力' + str(Gravity + i_gravity)

            g = yuc.Map(aa, weizz, yuchuli_len + first_len)
            g.subtosub()
            g.stosline()
            g.virnode()
            zuobiao = g.virzuobiao
            plt.scatter(zuobiao[:, 0], zuobiao[:, 1])

            for i in range(len(g.map)):
                for j in range(len(g.map)):
                    if g.map[i][j] == 1 and i < j and i in g.mianline and j in g.mianline:
                        x = [zuobiao[i][0], zuobiao[j][0]]
                        y = [zuobiao[i][1], zuobiao[j][1]]
                        plt.plot(x, y)
            path = 'D:\\tuoputu\\{0}\\{1}\\{2}'.format(time.strftime('%m.%d'), '初始长' + str(yuchuli_len + first_len),
                                                       '重力' + str(Gravity + i_gravity))
            folder = os.path.exists(path)
            if not folder:
                os.makedirs(path)
            plt.savefig('{0}\拓扑图{1}.png'.format(path, '开始'))
            plt.close()
            # print('这里这里', g.virzuobiao)
            Loss_out[env_value_name] = []
            for ir in range(50):
                # print(g.map, list(set(g.mianline) ^ set(g.bdz)), g.mianline)
                # a = gra.Gravity(g, 600, list(set(g.mianline) ^ set(g.bdz)), g.mianline)
                a = gra.Gravity(g, 10, g.mianline, g.mianline + [52, 53, 54, 55, 56], ir, Gravity + i_gravity)
                # a = gra.Gravity(g, ir * 10, g.mianline, g.mianline, ir)
                # all = [i for i in range(g.nodenum)]
                # a = gra.Gravity(g, 300, list(set(all) ^ set(g.bdz)), all)
                zuobiao = a
                coss = lossf.loss_function(g, g.mianline)
                # print(zuobiao)
                plt.scatter(zuobiao[:, 0], zuobiao[:, 1])
                for i in range(len(g.map)):
                    for j in range(len(g.map)):
                        if g.map[i][j] == 1 and i < j and i in g.mianline and j in g.mianline:
                            x = [zuobiao[i][0], zuobiao[j][0]]
                            y = [zuobiao[i][1], zuobiao[j][1]]
                            plt.plot(x, y)
                three_loss_weight = [coss[0], round(coss[1], 4), round(coss[4], 4), round(coss[5], 4)]
                # 全局

                loss_out = (coss[0] + 0.1) * (coss[1] + coss[4] + coss[5])
                Loss_out[env_value_name].append(loss_out)

                plt.title('{}'.format(three_loss_weight))
                plt.xlabel(str(round(coss[-1], 4))+' '+str(round(coss[-2], 4)))
                plt.axis('equal')
                path = 'D:\\tuoputu\\{0}\\{1}\\{2}'.format(time.strftime('%m.%d'), '初始长' + str(yuchuli_len + first_len),
                                                           '重力' + str(Gravity + i_gravity))
                folder = os.path.exists(path)
                if not folder:
                    os.makedirs(path)
                plt.savefig('{0}\拓扑图{1}.png'.format(path, ir))
                plt.close()
                # print('我日你大坝', coss)
            x = [i for i in range(len(Loss_out[env_value_name]))]

            # print(Loss_out)
            plt.plot(x, Loss_out[env_value_name])
            # plt.show()

        loss_folder = os.path.exists('D:\\tuoputu\\{0}\\LOSS'.format(time.strftime('%m.%d')))
        if not loss_folder:
            os.makedirs('D:\\tuoputu\\{0}\\LOSS'.format(time.strftime('%m.%d')))
        loss_path = 'D:\\tuoputu\\{0}\\LOSS\\LOSS{1}.csv'.format(time.strftime('%m.%d'),
                                                                 '初始长' + str(yuchuli_len + first_len))
        loss_data = pd.DataFrame(Loss_out)
        loss_data.to_csv(loss_path, index=False)


if __name__ == '__main__':
    # main(41, 391)
    aa = nclt.main(1)
    weizz = nclt.main(0)
    Loss_out = {}
    Loss_out['1580'] = []
    chushichangdu = 15
    zhongli = 170
    g = yuc.Map(aa, weizz, chushichangdu)
    g.subtosub()
    g.stosline()
    g.virnode()
    for ir in range(50):
        # print(g.map, list(set(g.mianline) ^ set(g.bdz)), g.mianline)
        # a = gra.Gravity(g, 600, list(set(g.mianline) ^ set(g.bdz)), g.mianline)
        a = gra.Gravity(g, 10, g.mianline, g.mianline + [52, 53, 54, 55, 56], ir, zhongli)
        # a = gra.Gravity(g, ir * 10, g.mianline, g.mianline, ir)
        # all = [i for i in range(g.nodenum)]
        # a = gra.Gravity(g, 300, list(set(all) ^ set(g.bdz)), all)
        zuobiao = a
        print('~~~~~~~~~~~~~~~~~~~第' + str(ir) + '次~~~~~~~~~~~~~~~~~~~~~~~~~~')
        coss = lossf.loss_function(g, g.mianline)
        # print(zuobiao)
        print(len(coss[-3]))
        print(coss[-3])
        plt.scatter(zuobiao[:, 0], zuobiao[:, 1])
        for i in range(len(g.map)):
            for j in range(len(g.map)):
                if g.map[i][j] == 1 and i < j and i in g.mianline and j in g.mianline:
                    x = [zuobiao[i][0], zuobiao[j][0]]
                    y = [zuobiao[i][1], zuobiao[j][1]]
                    plt.plot(x, y)
        three_loss_weight = [coss[0], round(coss[1], 4), round(coss[4], 4), round(coss[5], 4)]
        # 全局

        loss_out = (coss[0] + 0.1) * (coss[1] + coss[4] + coss[5])
        Loss_out['1580'].append(loss_out)

        plt.title('{}'.format(three_loss_weight))
        plt.xlabel(str(round(coss[-1], 4))+' '+str(round(coss[-2], 4)))
        plt.axis('equal')
        path = 'D:\\tuoputu\\test{0}\\{1}\\{2}'.format(time.strftime('%m.%d'), '初始长' + str(chushichangdu),
                                                       '重力' + str(zhongli))
        folder = os.path.exists(path)
        if not folder:
            os.makedirs(path)
        plt.savefig('{0}\拓扑图{1}.png'.format(path, ir))
        plt.close()
        # print('我日你大坝', coss)
    x = [i for i in range(len(Loss_out['1580']))]

    plt.plot(x, Loss_out['1580'])
    plt.show()

    loss_folder = os.path.exists('D:\\tuoputu\\test{0}\\LOSS'.format(time.strftime('%m.%d')))
    if not loss_folder:
        os.makedirs('D:\\tuoputu\\test{0}\\LOSS'.format(time.strftime('%m.%d')))
    loss_path = 'D:\\tuoputu\\test{0}\\LOSS\\LOSS{1}.csv'.format(time.strftime('%m.%d'),
                                                                 '初始长' + str(chushichangdu))
    loss_data = pd.DataFrame(Loss_out)
    loss_data.to_csv(loss_path, index=False)
